#!/usr/bin/python
'''
Created on Feb 10, 2012

@author: xzhu
'''

import sys

from operator import attrgetter

import stateHead
import allState, ioniState, hPath, const


def sort_path(lt):
    '''
    sort all the paths from low energy to high
    '''
    
    
    sorted_paths = []
    
    
    while lt:
        min_id = 0
        min_e = lt[0].highEner
        
        for i in range(len(lt)):
            if lt[i].highEner < min_e:
                min_id = i
                min_e = lt[i].highEner
                
        sorted_paths.append(lt.pop(min_id))
    
    return sorted_paths


def path_to_state(paths):
    '''
    Get all the actually existing states from paths and renumber the tier of states
    Return the states sorted by energy
    '''
    
    # all the states shown in the paths
    all_state = []
    
    
    count = 1
    
    # assign a tier to the states of a path
    for each_state in paths[0].states:
        each_state.tier = count
        count += 1
        all_state.append(each_state)
        
    # iterate through all the states of all the paths except the first path
    for i in range(len(paths)-1):
        for j in range(len(paths[i+1].states)):
            already_in = False
            
            # check whether a state of path is already one of the all_state(s)
            for in_state in all_state:
                if paths[i+1].states[j].is_same(in_state):
                    already_in = True
                    break
            
            # if not, then put that state into all_state and assign a tier for the state    
            if not already_in:
                paths[i+1].states[j].tier = j + 1
                all_state.append(paths[i+1].states[j])
    
    # how many layer are there on the path
    n_tier = len(paths[0].states)
    #print 'n_tire', n_tier
       
    # "return_states" are the sorted states, with "seq" attribute assigned
    return_states = []
    
    
    count = 1
    
    # assign the "seq" attribute for each state, which is the sequence number of a state
    for i in range(n_tier):
        i_tier_states = []
        
        for each_state in all_state:
            if each_state.tier == i+1:
                each_state.seq = count
                count += 1
                i_tier_states.append(each_state)
        
        # sort the states in one layer by their energies        
        sorted_tier = sorted(i_tier_states, key=attrgetter('ener'))
        
        # reassign the "seq" attribute after the states are sorted
        for ist in range(len(sorted_tier)):
            sorted_tier[ist].seq = count - len(sorted_tier) + ist
            return_states.append(sorted_tier[ist])
         
                 
    return return_states


def get_connects(paths, states):
    '''
    get the pairwise state to state transfer interconnection
    '''
    
    # all the possible connections between two states
    all_conn = []
    
    
    for each_path in paths:
        for i in range(len(each_path.states) - 1):
            first_seq = 0
            second_seq = 0
            
            
            for use_state in states:
                if each_path.states[i].is_same(use_state):
                    first_seq = use_state.seq
                    first_ener = use_state.ener
                    break
                
            for use_state in states:
                if each_path.states[i+1].is_same(use_state):
                    second_seq = use_state.seq
                    second_ener = use_state.ener
                    break
            
            ener_diff = second_ener - first_ener
            all_conn.append([first_seq, second_seq, ener_diff])
            
    uniq_conn = []
    uniq_conn.append(all_conn[0])
    
    for each_conn in all_conn:
        if each_conn not in uniq_conn:
            uniq_conn.append(each_conn)
            
    return uniq_conn

def update_state_count(states, paths):
    '''
    count for each state how many times it appears in paths
    '''
    for each_state in states:
        each_state.count = 0
        
    for each_path in paths:
        for each_state in each_path.states:
            for in_state in states:
                if each_state.is_same(in_state):
                    in_state.count += 1
                    break

def display_paths(paths):
    '''
    show all the paths
    the first row is the "seq" of all the intermediate states in the path,
        followed by the energy barrier (highest - initial) of this path.
    the following rows show the ionization states of the intermediates in the path
    '''
    
    fpath = open("paths.txt", 'w')
    for path in paths:
        for i in range(len(path.states)):
            print path.states[i].seq, '\t',
            fpath.write(str(path.states[i].seq) + '\t')
            
        #print path.highEner, path.states[0].ener, path.highEner - path.states[0].ener
        path.high_ener_states()
        print path.highEner - path.states[0].ener
        fpath.write(str(path.highEner - path.states[0].ener) + '\n')
        
        for i in range(len(path.states)):
            print path.states[i]
            fpath.write(str(path.states[i]) + '\n')
            
        fpath.write('\n')
        
    fpath.close()
        
        
def write_id_ener(seq, fname = 'id_ener.txt'):
    
    f_id_ener = open(fname, 'w')
    
    for i in range(len(seq)):
        for each_state in seq:
            if each_state.seq != i+1: continue
            f_id_ener.write('%s\t' % each_state.seq)
            print each_state.seq, '\t', 
            f_id_ener.write('%s\t' % each_state.ener)
            for crg in each_state.crgs:
                print const.INTTOSYMBOL[crg] + '\t',
                f_id_ener.write('%s\t' % str(const.INTTOSYMBOL[crg]))
            f_id_ener.write('\n')
            print each_state.ener, '\t', each_state.count 
    f_id_ener.close()
    
def write_id_ioni(seq, fname = 'id_ioni.txt'):
    
    f_id_ioni = open(fname, 'w')
    for each_conn in seq:
        print each_conn[0], '\t', each_conn[1], '\t', each_conn[2]
        f_id_ioni.write('%s\t%s\t%.3f\n' % (each_conn[0], each_conn[1], each_conn[2]))
        
                                  
def getBarrier(enerFile, initState):
    # stateLines stores all the lines from a file having the information of all the protonation states.    
    stateLines = open(enerFile).readlines()
    
    # stateHeader object stores information about the key residues.
    stateHeader = stateHead.KeyResidues()
    
    # use the first line of stateLines to initialize the residue header
    allKeyResidues = stateLines[0].split()[:-4]
    stateHeader.init_from_residue_names(allKeyResidues)
    
    # don't need the first line any more
    stateLines.pop(0)
    
    
    # set the header of state and load all the different protonation states
    all_poss_states = allState.AllState()
    #allState.AllState.sHead = stateHeader             # shallow copy, sHead is just a reference to stateHeader   
    #ioniState.IoniState.sHead = stateHeader           # shallow copy, sHead is just a reference to stateHeader 
    #hPath.HPath.sHead = stateHeader                   # shallow copy, sHead is just a reference to stateHeader
    # should use a class to store this stateHeader constant
    
    all_poss_states.load_all_states(stateLines) 
   
    
    
    # load the initial ionization state from the second file, like '0 0 0 + 0 -', it should be just one line
    initial_state = ioniState.IoniState()
    
    initialCharges = [const.SYMBOLTOINT[eachCharge] for eachCharge in open(initState).readline().split()]
    initial_state.load_init_state(all_poss_states, initialCharges)
    
    # get the first path from the initial state
    initial_path = hPath.HPath()
    initial_path.get_first_state(initial_state)
        
    
    # a queue of all the paths, run until all the paths in the queue are done
    path_queue = []
    path_queue.append(initial_path)     # now the path queue only has the initial path
    
    while not path_queue[0].is_done():
        new_path = path_queue.pop(0)
        
        # Get all the new paths which could be generated from this path
        buff_paths = new_path.next_step(all_poss_states)
        
        # Put all the new paths into the path queue
        path_queue.extend(buff_paths)
        
            
    #print 'There are %d different paths' % len(path_queue)
    #for each_path in path_queue:
    #    each_path.high_ener_states()
        
    #sorted_paths = sort_path(path_queue)   
    
    # sort all the paths by its highest energy among the states in it
    sorted_paths = sorted(path_queue, key=attrgetter('highEner')) 
    
    # all the intermediate states actually used in the paths 
    used_states = path_to_state(sorted_paths)
    update_state_count(used_states, sorted_paths)
    
    display_paths(sorted_paths)
    
    # output all the used ionization states in file "id_ener.txt" by default
    write_id_ener(used_states)
    
    # output all the connections between two states in file "id_ioni.txt" by default   
    all_conn = get_connects(sorted_paths, used_states)
    write_id_ioni(all_conn)
    
    
                       
def helpMessage():
    '''
    help message
    '''
    
    print "syntax: hBondPath.py -e ener.txt(all states, output of fix_io.py) init.txt(initial state)"
    print "        hBondPath.py -s pathinfo.txt"
    
    
def possible_protonation_states(pathFileName):
    
    sHead = stateHead.KeyResidues()
    fp = open(pathFileName, 'r')
    sHead.init_from_path_file(fp)
    
    all_poss_states = allState.AllState()
    all_poss_states.getAllPossibleStates()
    
    initial_state = ioniState.IoniState()
    initial_state.load_init_state(all_poss_states, sHead.initCharges)
    
    initial_path = hPath.HPath()
    initial_path.get_first_state(initial_state)
    
    path_queue = []
    path_queue.append(initial_path)     # now the path queue only has the initial path
    
    while not path_queue[0].is_done():
        new_path = path_queue.pop(0)
        
        # Get all the new paths which could be generated from this path
        buff_paths = new_path.next_step(all_poss_states)
        
        # Put all the new paths into the path queue
        path_queue.extend(buff_paths)
        
            
    print 'There are %d different paths' % len(path_queue)
    
    
    # all the intermediate states actually used in the paths 
    used_states = path_to_state(path_queue)
    
    out_fp = open("all_states.txt", 'w')
    
        
    for eachState in used_states:
        for eachCrg in eachState.crgs[:-1]:
            out_fp.write(const.INTTOSYMBOL[eachCrg] + '\t')
        out_fp.write(const.INTTOSYMBOL[eachState.crgs[-1]] + '\n')
    
    out_fp.close()

    
if __name__ == '__main__':
    
    # print the help message, if no argument or first option is '-h'.
    if len(sys.argv) == 1 or sys.argv[1] == "-h": 
        helpMessage()
        
    # -s: print all the different possible protonation states and save them in a file.
    elif sys.argv[1] == "-s":
        possible_protonation_states(sys.argv[2])
        
    # -e: get the energy barriers from the energies of all the possible states.
    elif sys.argv[1] == "-e":
        try:
            getBarrier(sys.argv[2], sys.argv[3])
        except AttributeError:
            import os
            sys.stderr.write("ERROR: " + os.getcwd() + '\n')
        